classdef BlochSimulator < handle
    %BLOCHSIMULATOR Solver for Bloch Equiation with Relaxation Terms
    %   Example:
    % blh=Simulation.BlochSimulator.BlochSimulator('Rp', 2*pi*1e3, 'OmegaY', 0, 'OmegaX', @(t) 2*pi*4.7e3*cos(2*pi*70e3*t), 'OmegaZ', 2*pi*70e3);
    % tic; [mVect, tList] = blh.solveTransient([0,1e-3], [0; 0; 0.5]);toc
    % blh.plot();
    %
    %=================================================================================================
    % The form of Bloch Equation this class used is:
    %
    %  dM                          Mxy       Mz                                              
    % ---- = - Omega \times M  -  -----  -  ---- + Rp * (M0 - M) ,    Mxy = M - Mz ,   Mz = (M \cdot ez)ez
    %  dt                          T2        T1                                                 
    %
    % Omega, M, Mxy, Mz, M0, ez are vectors; t, gamma, T1, T2, Rp are scalars. 
    % Mxy is the projection of M on the x-y plane, Mz is the z component of M. 
    % Omega is the angular velocity of Larmor Precession, Rp is the optical pumping rate, T1 and T2 are relaxation times.
    % M0 is the polarization of pump light, and it determines the steady-state polarization of M when relaxation and transverse field are both zero.
    %
    %  !!!!! Since we use a simplified form of Bloch Eq., the direction of main magnetic field should always be z-axis. !!!!!!!
    %
    %==================================================================================================

    
    properties(SetAccess = private)
        parameters
        solution
    end
    
    methods(Static)
        % auxiliary functions
        squFun = continuousSquareWave(t, freq, duty, rampRate);
        stepFun = continuousStepFunction(t, rampRate);
    end
    methods
        % solver
        [M, tList] = solveTransient(obj, tList, initM_vect);
        [M, tList] = solveSteadyState(obj, tRange, npts);
        % get data
        timeseriesSignal = getTimeSeriesSignal(obj, name, varargin);
        % show M motion
        showAnimateOfSolution(obj, varargin);
    end
    
    methods
        function obj = BlochSimulator(varargin)
            %BLOCHSIMULATOR Initialize parameters and check their format.
            %   arguments:
            %   'Rp': pumping rate [rad/s]
            %   'M0': photon polarization, 3D column vector, |M0|=0.5 for circularly pol. pump light; |M0|=0 for linearly pol.
            %         M0 = [0; 0; 0.5];    M0 = @(t)[0; 0; 0.5*(t>0).*(t<1)];
            %   'T1': T1 time in dark [s]
            %   'T2': T2 time in dark [s]
            %   'OmegaX', 'OmegaY', 'OmegaZ': 
            %          Angular velocity of Larmor Precession, should be a constant or function of time [rad/s]
            %
            %   'maxStep': max time step, default 100ns
            p = inputParser;
            paramCheck = @(x) isa(x, 'function_handle') || isnumeric(x);
            
            p.addParameter('Rp', 2*pi*1e3, paramCheck); 
            p.addParameter('M0', [0;0;0.5], paramCheck);     
            p.addParameter('T1', 1e-3, paramCheck);    
            p.addParameter('T2', 150e-6, paramCheck);   
            p.addParameter('OmegaX', 0.0, paramCheck); 
            p.addParameter('OmegaY', 0.0, paramCheck); 
            p.addParameter('OmegaZ', 2*pi*100e3, paramCheck); 
            p.addParameter('maxStep', 1e-6, @(x)isnumeric(x) && isscalar(x) && (x>0) ); 
            p.parse(varargin{:});
            
            param.Rp = obj.makeFunc(p.Results.Rp);
            param.M0 = obj.makeFunc(p.Results.M0);
            param.T1 = obj.makeFunc(p.Results.T1);
            param.T2 = obj.makeFunc(p.Results.T2);
            param.OmegaX = obj.makeFunc(p.Results.OmegaX);
            param.OmegaY = obj.makeFunc(p.Results.OmegaY);
            param.OmegaZ = obj.makeFunc(p.Results.OmegaZ);
            
            param.maxStep = p.Results.maxStep;
            
            % check parameter shape
            p2 = inputParser;
            
            scalarChecker = @(x)validateattributes(x(0),{'numeric'},{'scalar','real','nonnan','finite'});
            vector3Dchecker = @(x)validateattributes(x(0),{'numeric'},{'column','numel',3,'real','nonnan','finite'});
            p2.addRequired('Rp',scalarChecker);
            p2.addRequired('M0',vector3Dchecker);
            p2.addRequired('T1',scalarChecker);
            p2.addRequired('T2',scalarChecker);
            p2.addRequired('OmegaX',scalarChecker);
            p2.addRequired('OmegaY',scalarChecker);
            p2.addRequired('OmegaZ',scalarChecker);
            
            p2.parse(param.Rp, param.M0, param.T1, param.T2, param.OmegaX, param.OmegaY, param.OmegaZ);
            
            obj.parameters = param; 
        end
        
        %% plotter
        function plotSolution(obj)
            data = obj.getTimeSeriesSignal('solution');
            data.plot();
        end
        
        function plotM0(obj, tRange)
            if nargin < 2
                data = obj.getTimeSeriesSignal('M0');
            else
                data = obj.getTimeSeriesSignal('M0', tRange);
            end
            data.plot();
        end
        
        %% IO interface
        function M0=getM0(obj, tList)
            if nargin == 1
                M0 = obj.parameters.M0;
            else
                p = inputParser;
                p.addRequired('tList', @(t)validateattributes(t,{'numeric'},{'row','nonnan','finite'},'getM0', 'tList'));
                p.parse(tList);
                M0 = arrayfun(obj.parameters.M0, tList, 'UniformOutput', false);
                M0 = cell2mat(M0);
            end
        end

        function obj=setMaxStep(obj, maxStep)
            p = inputParser;
            p.addOptional('maxStep', 1e-6, @(x)isnumeric(x) && isscalar(x) && (x>0) ); 
            if nargin<2
                p.parse();
            else
                p.parse(maxStep);
            end
            obj.parameters.maxStep = p.Results.maxStep;
        end

    end

    
    methods(Static)

        function [demodX, demodY, demodR, demodTheta] = getDemod(tList, signal, refFreq, refPhase, k)
            if nargin == 5
                k = 1.0;
            end
            
            if length(tList) ~= length(signal) || length(tList) < 3
                error('signal must have the same length of tList, and length(tList) must greater than 2.');
            end
            
            dt = tList(2) - tList(1);
            period = 1.0/ refFreq / k; 
            refC = sqrt(2.0) * cos(k * 2*pi*refFreq*tList+refPhase);
            refS = sqrt(2.0) * sin(k * 2*pi*refFreq*tList+refPhase);            
            demodX = refC*signal*dt * 2.0/period;
            demodY = refS*signal*dt * 2.0/period;
            demodR = sqrt(demodX.*demodX + demodY.*demodY);
            demodTheta = atan(demodY./demodX);
        end
    end
    
    
    %% private methods
    methods (Static, Access = private)
        % Calculate the dM/dt vector
        % M and M0 should be function-handles that return 3D column vectors
        % T1, T2, Rp, OmegaX, OmegaY and OmegaZ should be function-handles that return scalars.
        function dMdt = blochEq(t, M, T1, T2, OmegaX, OmegaY, OmegaZ, Rp, M0)
            g1t = 1./T1(t) + Rp(t);
            g2t = 1./T2(t) + Rp(t);
            oXt = OmegaX(t);
            oYt = OmegaY(t);
            oZt = OmegaZ(t);
            RpVect = Rp(t)*M0(t);
            
            A = [-g2t, oZt, -oYt; ...
                 -oZt, -g2t, oXt; ...
                 oYt, -oXt, -g1t];
            dMdt = A*M + RpVect;
        end
        
        function funcHdl = makeFunc(param)
            if isnumeric(param)
                funcHdl = @(t) param;
            elseif isa(param, 'function_handle')
                funcHdl = param;
            else
                error('Wrong type of param');
            end
        end
        
        function res = periodCheck(func, t0, t1, tw)
            if nargin < 4
                tw = abs(t1-t0);
            end
            npts = 521;
            dt = linspace(-tw/2,+tw/2,npts);
            try
                f0 = func(t0+dt);
                f1 = func(t1+dt);
            catch
                f0 = arrayfun(func, t0+dt, 'UniformOutput', false); f0 = cell2mat(f0);
                f1 = arrayfun(func, t1+dt, 'UniformOutput', false); f1 = cell2mat(f1);
            end
            s = sum(sum(abs(f0) + abs(f1)))+eps;
            res = sum(sum(abs(f0 - f1)))/s;
        end
    end
end

